from sklearn.datasets import load_iris
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import cross_val_score  # 引入 K 折交叉验证

iris = load_iris()
target_names = iris.target_names


def knn_cross_val_score(n_neighbors, features=iris.data, labels=iris.target):
    knn_classifier = KNeighborsClassifier(n_neighbors=n_neighbors)
    scores = cross_val_score(knn_classifier, features, labels, cv=10)
    return scores.mean()


def predict_iris_species_with_knn_classifier(n_neighbors, test_features, features=iris.data, labels=iris.target):
    knn_classifier = KNeighborsClassifier(n_neighbors=n_neighbors)
    knn_classifier.fit(features, labels)
    predict_labels = knn_classifier.predict(test_features)
    predict_species = [target_names[label] for label in predict_labels]
    return predict_species


optimal_k = 5
for k in range(5, 11):
    print(f"k={k}: {knn_cross_val_score(k)}")
    optimal_k = k if knn_cross_val_score(k) > knn_cross_val_score(optimal_k) else optimal_k
print(f"optimal_k: {optimal_k}")
print(predict_iris_species_with_knn_classifier(optimal_k, [[1.5, 3, 5.8, 2.2], [6.2, 2.9, 4.3, 1.3]]))
